
# --*-- coding:utf-8 --*--
"""
    - 能够同时抽取文件中的的文字和图片
    - pdf 和 docx(不支持doc)，可以根据平台进行转换
    - 图片和文本暂时存储到一个临时的文件夹，看师兄如何调用image caption ， 后续应该删除，没有必要存在本地
    - docx支持的图片格式可以通过规则增加
    - 更换了解析pdf的库，但是还没有自己测试速度
"""

# from UltraSearch.im2txt.extract_info import extract

from pyhanlp import *
from abstract.text import parse_pdf , parse_document , parse_txt
from abstract.image import extract_doc , extract_pdf


def load_stopwords(path):
    try:
        with open(path,encoding='utf-8') as f:
            words = f.read().split('\n')
        f.close()
    except:
        # 
        return ['a','the']
    return words

def filter_keys(keywords:str , stopwords:str):
    """
        定义好了停词列表
    """
    keywords = keywords.split('##')
    new_keys = []
    for key in keywords:
        if key not in stopwords:
            new_keys.append(key)
    return "##".join(new_keys)

def extract_keys(document,keys_count = 3):
    """
        also for image captions , 抽取3关键字
    """
    try:
        # a = 1/0
        TextRankKeyword = JClass("com.hankcs.hanlp.summary.TextRankKeyword")
        keywords = HanLP.extractKeyword(document, keys_count)
    except:
        return "skip this file!"

    return '##'.join(keywords) 
    

def extract_sents(document , sentences_count = 1):
    """
        textcnn抽取文档的关键字 ， 抽取一句
    """
    try:
        TextRankSentence = JClass("com.hankcs.hanlp.summary.TextRankSentence")
        sentences = HanLP.extractSummary(document, sentences_count)
    except:
        return "skip this file!"

    return '##'.join(sentences) 


def get_image_path(dirs):
    """
        每张图片的绝对路径以列表形式返回
    """
    import os 
    paths =[]
    for root , _ , files in  os.walk(dirs):
        for file in files:
            paths.append(os.path.join(root,file))
    return paths


def get_image(file,tmp_dir):
    """
        将包含图片的文件图片抽取到临时文件夹【只支持pdf和doc】
        并且将每张临时图片的路径以字符串返回[使用换行符作为分割]
    """

    import os 
    file_type = os.path.splitext(file)[-1]

    if file_type =='.pdf':
        extract_pdf(file , tmp_dir)
        path = get_image_path(tmp_dir)

    elif file_type =='.docx':
        extract_doc(file , tmp_dir)
        path = get_image_path(tmp_dir)
    return path

def abstract_file(file):
    """
        todo: some bugs to fix 
        return file abstract and keys
    """
    import os
    file_type = os.path.splitext(file)[-1]

    if file_type == '.docx':
        document  = parse_document(file)
    elif file_type == '.pdf':
        document = parse_pdf(file)
    elif file_type == '.txt':
        document = parse_txt(file) 

    
    keys = extract_keys(document)
    sents = extract_sents(document)

    return keys , sents
# def translate_zh(sentence):
#     """
#         将英文标题翻译成中文并且分词，返回同样的keys ，sents
#         如果有网则进行翻译，否则只能够输出英文的摘要和关键字
#     """

#     #  判断是否联网 ， 如果为0则是有网
#     exit_code = os.system('ping translate.google.cn')

#     if exit_code==0:
#         from googletrans import Translator

#         translator = Translator(service_urls=['translate.google.cn'], proxies={'http': "localhost:80"})
#         chinese_sentence = translator.translate(sentence, src='en', dest='zh-cn').text

#         keys_zh = extract_keys(chinese_sentence )
#         sents_zh = chinese_sentence

#         keys = keys_zh 
#         sents = sents_zh

#     else:
#         keys = extract_keys(sentence)
#         sents = sentence

#     return keys , sents
    






# if __name__ == '__main__':

    # tmp_dir ='./ts

    # clear_temp(tmp_dir)

    # stopwords = load_stopwords('abstract/stopwords.txt')
    

    # ==========clear tmp dir : socket flag
    # clear_temp(tmp_dir=tmp_dir)

    # ==== word  ===
    # file_path = 'Abst-example/doc/航发黎阳招聘简章.docx'
    # file_path = 'Abst-example/image/image.docx'
    # file_path = 'Abst-example/doc/假期须知.docx'

    # def test_doc(file_path , tmp_dir = './temp'):
    #     extract_doc(file_path , tmp_dir = tmp_dir)
    #     path = get_image_path(file_path , dirs=tmp_dir)
    #     return tmp_dir , path
    # tmp_dir ,path= test_pdf(file_path , tmp_dir='./tmp')
    # #==========clear tmp dir : socket flag
    # clear_temp(tmp_dir='./temp')


    
   


# method 2
#  python -m fitz extract  -images 'Abst-example/pdf/2021.naacl-main.118.pdf' -output ./test
